
from models import *

from torch.fx import GraphModule, symbolic_trace
from torch.fx.node import Node
import torch.nn as nn

def collect_nodes_for_observer(model: GraphModule):
    """Collects the nodes where input and output quantization should be added."""
    input_nodes, output_nodes = [], []

    skip_op_types = {nn.AvgPool2d, nn.ReLU, nn.Dropout}         # Define the types to skip

    for node in model.graph.nodes:
        if node.op == 'call_module':
            if type(dict(model.named_modules())[node.target]) in skip_op_types:
                continue
            # Collect input nodes
            if isinstance(node.args[0], Node):
                input_nodes.append((node.args, node))
            # Collect output nodes
            for user in node.users:
                output_nodes.append((node, user))
    return input_nodes, output_nodes

def insert_observer_nodes(model: nn.Module, observer) -> GraphModule:
    """
    Insert input and output observer nodes into the model.

    Args:
        model (nn.Module): The model to modify.
        observer: The observer class to instantiate (e.g., torch.quantization.HistogramObserver).

    Returns:
        GraphModule: The modified model with observer nodes inserted.
    """
    traced_model = symbolic_trace(model)
    graph = traced_model.graph

    input_nodes, output_nodes = collect_nodes_for_observer(traced_model)

    # Insert input observer nodes
    for src_node, target_node in input_nodes:

        observer_name = f'{target_node}_input_observer'
        traced_model.add_module(observer_name, copy.deepcopy(observer))
        with graph.inserting_before(target_node):
            input_observer_node = graph.create_node('call_module', observer_name, (src_node[0],), {})
            target_node.replace_input_with(src_node[0], input_observer_node)

    # Insert output observer nodes
    for src_node, user_node in output_nodes:
        observer_name = f'{src_node}_output_observer'
        traced_model.add_module(observer_name, copy.deepcopy(observer))
        with graph.inserting_after(src_node):
            output_observer_node = graph.create_node('call_module', observer_name, (src_node,), {})
            user_node.replace_input_with(src_node, output_observer_node)
    
    graph.lint()
    traced_model.recompile()
    return traced_model

if __name__ == '__main__':
    # Example usage
    observer=torch.quantization.observer.HistogramObserver(
        dtype=torch.qint8, qscheme=torch.per_tensor_affine,
        quant_min=-128, quant_max=127)
    model = ResNet18()
    model.eval()
    quantized_model = insert_observer_nodes(model, observer)
    quantized_model.graph.print_tabular()



